package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/coder/websocket"
)

// Message 消息结构，需要与服务端保持一致
type Message struct {
	Type    string `json:"type"`
	Content string `json:"content,omitempty"`
	Topic   string `json:"topic,omitempty"`
}

func main() {
	if len(os.Args) < 2 {
		log.Fatal("Usage: go run client.go <uid>")
	}
	uid := os.Args[1]

	// 连接到WebSocket服务器
	ctx, cancel := context.WithTimeout(context.Background(), time.Second*10)
	defer cancel()

	conn, _, err := websocket.Dial(ctx, fmt.Sprintf("ws://localhost:8081/ws?uid=%s", uid), nil)
	if err != nil {
		log.Fatalf("Failed to connect to WebSocket server: %v", err)
	}
	defer conn.Close(websocket.StatusNormalClosure, "client closing connection")

	log.Printf("Connected to WebSocket server with UID: %s", uid)

	// 订阅一个主题
	subscribeMsg := Message{
		Type:  "subscribe",
		Topic: "news",
	}
	if err := sendMessage(ctx, conn, subscribeMsg); err != nil {
		log.Fatalf("Failed to subscribe: %v", err)
	}
	log.Println("Subscribed to topic: news")

	// 捕获系统信号，优雅关闭
	stop := make(chan os.Signal, 1)
	signal.Notify(stop, os.Interrupt, syscall.SIGTERM)

	// 启动一个goroutine来接收消息
	done := make(chan struct{})
	go receiveMessages(conn, done)

	// 等待中断信号或接收goroutine结束
	select {
	case <-stop:
		log.Println("Received interrupt signal, closing connection...")
	case <-done:
		log.Println("Connection closed by server")
	}

	// 取消订阅
	unsubscribeMsg := Message{
		Type:  "unsubscribe",
		Topic: "news",
	}
	ctx, cancel = context.WithTimeout(context.Background(), time.Second*5)
	defer cancel()
	if err := sendMessage(ctx, conn, unsubscribeMsg); err != nil {
		log.Printf("Failed to unsubscribe: %v", err)
	} else {
		log.Println("Unsubscribed from topic: news")
	}
}

// sendMessage 发送消息到服务器
func sendMessage(ctx context.Context, conn *websocket.Conn, msg Message) error {
	data, err := json.Marshal(msg)
	if err != nil {
		return fmt.Errorf("failed to marshal message: %w", err)
	}

	if err := conn.Write(ctx, websocket.MessageText, data); err != nil {
		return fmt.Errorf("failed to write message: %w", err)
	}

	return nil
}

// receiveMessages 接收来自服务器的消息
func receiveMessages(conn *websocket.Conn, done chan struct{}) {
	defer close(done)

	for {
		_, data, err := conn.Read(context.Background())
		if err != nil {
			if websocket.CloseStatus(err) == websocket.StatusNormalClosure {
				log.Println("Connection closed normally")
			} else {
				log.Printf("Error reading from server: %v", err)
			}
			return
		}

		var msg Message
		if err := json.Unmarshal(data, &msg); err != nil {
			log.Printf("Failed to unmarshal message: %v", err)
			continue
		}

		log.Printf("Received message: %+v", msg)
	}
}
